{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5.3.1 Using PETSc \n",
    "We learn how we can interface the popular parallel toolkit PETSc for solving linear equations and much more. We use the Python interface petsc4py. The whole Python file is [n2p_ex1.py](n2p_ex1.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipyparallel import Cluster\n",
    "c = await Cluster(engines=\"mpi\").start_and_connect(n=4, activate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "from ngsolve import *\n",
    "from netgen.occ import unit_square\n",
    "from mpi4py.MPI import COMM_WORLD as comm\n",
    "\n",
    "ngmesh = unit_square.GenerateMesh(maxh=0.1, comm=comm)\n",
    "for l in range(4):\n",
    "    ngmesh.Refine()\n",
    "mesh = Mesh(ngmesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "import numpy as np\n",
    "import petsc4py.PETSc as psc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "fes = H1(mesh, order=1)\n",
    "u,v = fes.TnT()\n",
    "a = BilinearForm(grad(u)*grad(v)*dx+u*v*ds).Assemble()\n",
    "f = LinearForm(x*v*dx).Assemble()\n",
    "gfu = GridFunction(fes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We convert the local sparse matrix to a local PETSc AIJ matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "locmat = a.mat.local_mat\n",
    "val,col,ind = locmat.CSR()\n",
    "ind = np.array(ind, dtype='int32')\n",
    "apsc_loc = psc.Mat().createAIJ(size=(locmat.height, locmat.width), csr=(ind,col,val), comm=MPI.COMM_SELF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The NGSolve ParallelDofs object corresponds to a PETSc IndexSet object, which we create next. In PETSc, dofs are globally enumerated, what is not the case in NGSolve. For this purpose, a ParallelDofs class can generate a globally consistent enumeration of dofs. The generated globnums array contains the global dof-numbers of the local dofs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "pardofs = fes.ParallelDofs()\n",
    "globnums, nglob = pardofs.EnumerateGlobally() \n",
    "# print (list(globnums))\n",
    "\n",
    "iset = psc.IS().createGeneral (indices=globnums, comm=comm)\n",
    "lgmap = psc.LGMap().createIS(iset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create the global matrix using our local2global map:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "mat = psc.Mat().createPython(size=nglob, comm=comm)\n",
    "mat.setType(psc.Mat.Type.IS)\n",
    "mat.setLGMap(lgmap)\n",
    "mat.setISLocalMat(apsc_loc)\n",
    "mat.assemble()\n",
    "mat.convert(\"mpiaij\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "f.vec.Cumulate()\n",
    "\n",
    "v1, v2 = mat.createVecs()\n",
    "\n",
    "v2loc = v2.getSubVector(iset)\n",
    "v2loc.getArray()[:] = f.vec.FV()\n",
    "v2.restoreSubVector(iset, v2loc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "ksp = psc.KSP()\n",
    "ksp.create()\n",
    "ksp.setOperators(mat)\n",
    "ksp.setType(psc.KSP.Type.CG)\n",
    "ksp.setNormType(psc.KSP.NormType.NORM_NATURAL)\n",
    "ksp.getPC().setType(\"gamg\")\n",
    "# ksp.getPC().setGAMGType(psc.PC.GAMGType.CLASSICAL)\n",
    "# ksp.getPC().setGAMGLevels(10)\n",
    "ksp.setTolerances(rtol=1e-6, atol=0, divtol=1e16, max_it=4000)\n",
    "ksp.solve(v2,v1)   \n",
    "\n",
    "printmaster (\"petsc-its =\", ksp.its)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "v1loc = v1.getSubVector(iset)\n",
    "for i in range(len(gfu.vec)):\n",
    "    gfu.vec.FV()[i] = v1loc.getArray()[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from ngsolve.webgui import Draw\n",
    "gfu = c[:][\"gfu\"]\n",
    "Draw (gfu[0]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
